from llm_controller.llm_agent_action import *
from llm_controller.llm_agent_negotiation_system import *  # system perspective negotiation
from llm_controller.memory import DrivingMemory
import highway_env
import imageio
import openpyxl
import os
import shutil
import logging
from config import config, logger

# 设置日志
logger = logging.getLogger(__name__)

def open_excel(i):
    file_dir = './llm_controller/excel/' + '/'
    file_name = file_dir + str(i) + '.xlsx'

    if not os.path.exists(file_dir):
        os.makedirs(file_dir)
    workbook = openpyxl.Workbook()
    if os.path.exists(file_name):
        workbook = openpyxl.load_workbook(file_name)

    # if 'Sheet' in workbook.sheetnames:
    #     del workbook['Sheet']
    return file_name, workbook

def write_data(workbook, env, t):
    column_names = ['t', 'x', 'y', 'v', 'theta', 'background_veh?']
    for vehicle in env.road.vehicles:
        sheet_name = str(vehicle.id)
        if sheet_name not in workbook.sheetnames:
            worksheet = workbook.create_sheet(sheet_name)
            worksheet.append(column_names)
        else:
            worksheet = workbook[sheet_name]
        controlled_vehicles = env.controlled_vehicles
        if vehicle not in controlled_vehicles:
            background_vehicles = False
        else:
            background_vehicles = True
        state = [round(vehicle.position[0], 2), round(vehicle.position[1], 2), round(vehicle.speed, 2), round(vehicle.heading, 2), background_vehicles]
        row_data = [t, round(vehicle.position[0], 2), round(vehicle.position[1], 2), round(vehicle.speed, 2), round(vehicle.heading, 2), background_vehicles]
        worksheet.append(row_data)
        worksheet.cell(row=t + 2, column=1, value=t)
        for i, item in enumerate(state):
            worksheet.cell(row=t + 2, column=i + 2, value=item)
    return workbook

# if choose merge, active config
# config = {
#     "simulation_frequency": 20,
#     "policy_frequency": 5,
#     "duration":40
# }

# env = gym.make('merge-multi-agent-v0', config=config)
env = gym.make('intersection-multi-agent-v0')
# env = gym.make('highway-v0')


for i in range(100):
    try:
        logger.info(f"Starting simulation {i}")
        video_path = './llm_controller/video/' + str(i) + '.mp4'  
        writer = imageio.get_writer(video_path, fps=30) 
        file_name, workbook = open_excel(i)
        terminated = False
        t = 0
        obs = env.reset()
        
        while not (terminated):
            print('---------------------------------------------------------------')
            try:
                # memory module
                memory = DrivingMemory(env)

                # negotiation module
                llm_agent_conflict_resolver = LlmAgent_negotiation_module(env)  # system perspective negotiation
                negotiation_prompt, conflicting_info = llm_agent_conflict_resolver.llm_controller_run(env)

                # decision
                llm_agent = LlmAgent_action_module(env)  
                sce = llm_agent.retrun_sce()  # sce data
                llm_actions = llm_agent.llm_controller_run(env, negotiation_prompt, conflicting_info, env.controlled_vehicles, memory)  # negotiation results from upper layer and conflict info which stores distance speed

                action = [item for sublist in llm_actions for item in sublist]  # [[1], [3]]->[1, 3]

                for veh in env.controlled_vehicles:
                    print(veh, veh.speed)

                obs, global_reward, terminated, info = env.step(tuple(action), env)
                env.render()
                print("llm_actions:", action)
                print("global_reward_llm:", global_reward)

                frame = env.render('rgb_array')
                writer.append_data(frame)

                workbook = write_data(workbook, env, t)
                workbook.save(file_name)
                t += 1

            except Exception as e:
                logger.error(f"Error in simulation step {t}: {e}")
                # 使用默认动作继续
                action = [1] * len(env.controlled_vehicles)  # IDLE action
                obs, global_reward, terminated, info = env.step(tuple(action), env)
                env.render()
                t += 1
                
        writer.close()
        logger.info(f"Completed simulation {i}")
        print(i)
        
    except Exception as e:
        logger.error(f"Error in simulation {i}: {e}")
        print(f"Simulation {i} failed: {e}")
        continue

